#Bibliotecas usadas
import cv2
import numpy as np
import matplotlib.pyplot as plt
from PIL import Image
from pylab import *
from IPython.display import display, Image

#Caminho do arquivo
cap = cv2.VideoCapture('caleido.mp4')

#Condição para caso haja erro na leitura do arquivo
if not cap.isOpened():
    print('Erro na leitura')
    exit()

#Loop que vai imprimir os frames enquanto o vídeo não finalizar
while cap.isOpened():
    ret, frame = cap.read()
                
    # Converter RGB para HSV
    hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)
    
    #Define o pixel vermelho que deve ser detectado
    light_red = np.array([0,100,100])
    dark_red = np.array([10,255,255])
    mask1 = cv2.inRange(hsv, light_red, dark_red)

    #Será criada uma segunda máscara, para 
    light_red = np.array([170,100,100])
    dark_red = np.array([180,255,255])
    mask2 = cv2.inRange(hsv, light_red, dark_red)

    #Aplica a máscara à imagem HSV para captar apenas os pixels especificados
    mask = mask1 + mask2
    
    #Verifica cada pixel para confirmar se contém as cores especificadas
    output = cv2.bitwise_and(frame,frame, mask = mask)
    
    if ret:
        #Exibe os frames em ordem comparando ambos vídeos
        fig, ax = plt.subplots(1, 2)
        ax[0].imshow(frame[:, :, ::-1])
        ax[0].set_title('Vídeo Original')
        ax[1].imshow(output[:, :, ::-1])
        ax[1].set_title('Somente vermelho')
        plt.show()
        
    else:
        break

cap.release()
cv2.destroyAllWindows()

#Verificação do método com outro vídeo
cap = cv2.VideoCapture('mont_russa.mp4')

#Condição para caso haja erro na leitura do arquivo
if not cap.isOpened():
    print('Erro na leitura')
    exit()

#Loop que vai imprimir os frames enquanto o vídeo não finalizar
while cap.isOpened():
    ret, frame = cap.read()
                
    # Converter RGB para HSV
    hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)
    
    #Define o pixel vermelho que deve ser detectado
    light_red = np.array([0,100,100])
    dark_red = np.array([10,255,255])
    mask1 = cv2.inRange(hsv, light_red, dark_red)

    #Será criada uma segunda máscara, para 
    light_red = np.array([170,100,100])
    dark_red = np.array([180,255,255])
    mask2 = cv2.inRange(hsv, light_red, dark_red)

    #Aplica a máscara à imagem HSV para captar apenas os pixels especificados
    mask = mask1 + mask2
    
    #Verifica cada pixel para confirmar se contém as cores especificadas
    output = cv2.bitwise_and(frame,frame, mask = mask)
    
    if ret:
        #Exibe os frames em ordem comparando ambos vídeos
        fig, ax = plt.subplots(1, 2)
        ax[0].imshow(frame[:, :, ::-1])
        ax[0].set_title('Vídeo Original')
        ax[1].imshow(output[:, :, ::-1])
        ax[1].set_title('Somente vermelho')
        plt.show()
        
    else:
        break

cap.release()
cv2.destroyAllWindows()